#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
COCO to VOC
"""
import sys
import argparse
import json
import os
from xml.dom.minidom import Document


def generate_dict(json_file, category):
    """
    generate dict from json file
    :param json_file:
    :return:
    """

    dictionary = {}
    js_file = json.load(open(json_file))

    for img in js_file["images"]:
        if img['id'] not in dictionary:
            dictionary[img['id']] = {}
        dictionary[img['id']]['file_name'] = img['file_name']
        dictionary[img['id']]['height'] = img['height']
        dictionary[img['id']]['width'] = img['width']
        dictionary[img['id']]['boxs'] = []

    for anno in js_file["annotations"]:
        if anno['category_id'] in category:
            bbox = anno['bbox']
            bbox.append(anno['category_id'])  # add the object class id
            dictionary[anno['image_id']]['boxs'].append(bbox)

    return dictionary


def add_object(doc, annotation, obj_name,
               xmin_v, ymin_v, xmax_v, ymax_v, occl):
    """
    add xml object info
    :param doc: xml object
    :param annotation: annotation
    :param obj_name: class name
    :param xmin_v: left
    :param ymin_v:  top
    :param xmax_v:  right
    :param ymax_v: bottom
    :param occl: truncted of not
    :return: None
    """
    obj = doc.createElement("object")

    objname = doc.createElement("name")
    objname.appendChild(doc.createTextNode(obj_name))
    pose = doc.createElement("pose")
    pose.appendChild(doc.createTextNode("front"))
    truncated = doc.createElement("truncated")
    truncated.appendChild(doc.createTextNode(str(occl)))
    difficult = doc.createElement("difficult")
    difficult.appendChild(doc.createTextNode("0"))
    obj.appendChild(objname)
    obj.appendChild(pose)
    obj.appendChild(truncated)
    obj.appendChild(difficult)

    bndbox = doc.createElement("bndbox")
    xmin = doc.createElement("xmin")
    ymin = doc.createElement("ymin")
    xmax = doc.createElement("xmax")
    ymax = doc.createElement("ymax")
    xmin.appendChild(doc.createTextNode(str(xmin_v)))
    ymin.appendChild(doc.createTextNode(str(ymin_v)))
    xmax.appendChild(doc.createTextNode(str(xmax_v)))
    ymax.appendChild(doc.createTextNode(str(ymax_v)))
    bndbox.appendChild(xmin)
    bndbox.appendChild(ymin)
    bndbox.appendChild(xmax)
    bndbox.appendChild(ymax)
    obj.appendChild(bndbox)

    annotation.appendChild(obj)


def coco_to_voc(coco_dir, voc_dir, category):
    """
    convert coco labels to voc labels
    :param coco_dir: coco root directory
    :param voc_dir:  voc root directory
    :param category: a dict that represent voc
    class name exists in coco class. for example,
     class 'person' exists both in coco and voc, class idx in coco is '0'
    :return: None
    """
    json_file = os.path.join(coco_dir, "annotations", "instances_train.json")
    if not os.path.exists(json_file):
        exit("json file not found.")
    dictionary = generate_dict(json_file, category)
    img_count = len(dictionary)
    count = 0
    print("convert coco to voc ...")
    for img_id in dictionary:
        info = dictionary[img_id]
        count += 1
        if count % 100 == 0:
            print("Complete: {}, Total: {}".format(count, img_count))

        if len(info['boxs']) < 1:
            continue
        doc = Document()

        n_rows, n_cols, channels = info['height'], info['width'], 3

        annotation = doc.createElement('annotation')
        doc.appendChild(annotation)

        folder = doc.createElement("folder")
        folder.appendChild(doc.createTextNode("COCO"))
        annotation.appendChild(folder)

        filename = doc.createElement("filename")
        filename.appendChild(doc.createTextNode(info['file_name']))
        annotation.appendChild(filename)

        source = doc.createElement("source")
        database = doc.createElement("database")
        database.appendChild(doc.createTextNode("The COCO Database"))
        anno = doc.createElement("annotation")
        anno.appendChild(doc.createTextNode("HIKDATA_AIA COCO"))
        image = doc.createElement("image")
        image.appendChild(doc.createTextNode("flickr"))
        source.appendChild(database)
        source.appendChild(anno)
        source.appendChild(image)
        annotation.appendChild(source)

        sizeimage = doc.createElement("size")
        imagewidth = doc.createElement("width")
        imagewidth.appendChild(doc.createTextNode(str(n_cols)))
        imageheight = doc.createElement("height")
        imageheight.appendChild(doc.createTextNode(str(n_rows)))
        imagedepth = doc.createElement("depth")
        imagedepth.appendChild(doc.createTextNode(str(channels)))
        sizeimage.appendChild(imagedepth)
        sizeimage.appendChild(imagewidth)
        sizeimage.appendChild(imageheight)
        annotation.appendChild(sizeimage)
        for box in info['boxs']:
            obj_name = category[box[4]]
            xmin_v = int(float(box[0]))
            ymin_v = int(float(box[1]))
            xmax_v = int(float(box[0])+float(box[2]))
            ymax_v = int(float(box[1])+float(box[3]))
            add_object(doc, annotation, obj_name, xmin_v, ymin_v, xmax_v, ymax_v, 0)

        xml_name = info['file_name'].split('.')[0]+'.xml'
        xml_file = open(os.path.join(voc_dir, "Annotations", xml_name), 'w')
        xml_file.write(doc.toprettyxml(indent=' '*4))


if __name__ == '__main__':
    try:
        VOC_ROOT_DIR = sys.argv[1]
        COCO_SAVE_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--voc_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/voc2012"),
                            help="VOC directory")
        PARSER.add_argument("--coco_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/coco"),
                            help="COCO directory")

        ARGS = PARSER.parse_args()
        VOC_ROOT_DIR = ARGS.voc_dir
        COCO_SAVE_DIR = ARGS.coco_dir
    LABELS = ['person', 'bicycle', 'car', 'motorbike',
              'aeroplane', 'bus', 'train', 'truck', 'boat',
              'traffic light', 'fire hydrant', 'stop sign',
              'parking meter', 'bench', 'bird', 'cat', 'dog',
              'horse', 'sheep', 'cow', 'elephant', 'bear',
              'zebra', 'giraffe', 'backpack', 'umbrella',
              'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
              'snowboard', 'sports ball', 'kite', 'baseball bat',
              'baseball glove', 'skateboard', 'surfboard',
              'tennis racket', 'bottle', 'wine glass', 'cup',
              'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
              'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog',
              'pizza', 'donut', 'cake', 'chair', 'sofa', 'pottedplant',
              'bed', 'diningtable', 'toilet', 'tvmonitor', 'laptop',
              'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
              'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
              'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

    CATEGORIES = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorbike', 4: 'aeroplane',
                  5: 'bus', 6: 'train', 39: 'bottle', 8: 'boat', 14: 'bird',
                  15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow',
                  56: 'chair', 57: 'sofa', 58: 'pottedplant', 60: 'diningtable', 62: 'tvmonitor'}
    coco_to_voc(COCO_SAVE_DIR, VOC_ROOT_DIR, CATEGORIES)
